from ..base import MemoryProviderBase, logger
import time
import json
import os
import yaml
from config.config_loader import get_project_dir
from core.utils.inner_request_util import inner_send_mem
import threading


short_term_memory_prompt = """
你是一位专业的小朋友心理语言分析的专家，请根据输入的小朋友对话内容，生成一份两段式的情绪分析报告。注意语言需客观、简洁、中立、专业，结构清晰，通熟易懂。

## 分析一：要求100字左右，总结对话主题和高频词，并指出对话主题和高频词在对话中的具体位置。
请先描述本轮对话内容(如围绕某某主题展开)，详细分析小朋友在对话中的情绪变化，并提取出频率较高的情绪相关词汇(如“好烦”、“兴奋”、“困惑”、“我不喜欢”、“没人陪我玩”等)，同时指出它们在对话中大致的出现位置（例如第2轮、第5轮等），最后进行对话总结。
如在 xxx 小朋友的本轮对话，主要围绕【事件主题，如“与同伴的争执”或“家里的小宠物”】展开，期间多次表达出【核心情绪，如“烦躁”“兴奋”或“困惑”】，分别出现在对话的【第几轮】中。整体语言表现出【如“跳跃性思维”“情绪波动”或“好奇心强烈”等特点】。

## 分析二：要求150字左右，围绕“具体情绪特征+个性表现+建议”展开，深入分析小朋友在某轮对话（如第3轮、第6轮等）中表现出的情绪特征（如愤怒、焦虑、快乐等），结合语言内容识别其潜在的性格倾向（如乐观、敏感、自我调节力强等），并给予针对性的建议（如“可以通过绘本帮助理解情绪”“建议鼓励其表达感受”等）。注意语言要专业、温和、有引导性。
如在第 xx 轮中，xxx 小朋友提到【具体内容】时，表现出明显的【如“沮丧”“兴奋”或“内疚”等情绪】，同时使用了诸如【具体词汇】等表达，体现出其具有【如“敏感细腻”“责任心强”“缺乏安全感”等性格特征】。这可能与其【背景信息，如学校适应、人际关系等】有关。建议家长/老师可以通过【具体建议，如“角色扮演”“集体活动引导”等方式】，帮助其更好地识别与调节情绪，提升社交与表达能力。
如在第 xxx 轮中...

# 上面输出的分析一和分析二需要完整保留，放到result1和result2字段，并以可解析的json字符串返回
"""


def extract_json_data(json_code):
    start = json_code.find("```json")
    # 从start开始找到下一个```结束
    end = json_code.find("```", start + 1)
    # print("start:", start, "end:", end)
    if start == -1 or end == -1:
        try:
            jsonData = json.loads(json_code)
            return json_code
        except Exception as e:
            print("Error:", e)
        return ""
    jsonData = json_code[start + 7: end]
    return jsonData


TAG = __name__


class MemoryProvider(MemoryProviderBase):
    def __init__(self, config):
        super().__init__(config)
        self.short_momery = ""
        self.memory_path = get_project_dir() + "data/.memory.yaml"
        self.load_memory()

    def init_memory(self, role_id, llm):
        super().init_memory(role_id, llm)
        self.load_memory()

    def load_memory(self):
        all_memory = {}
        if os.path.exists(self.memory_path):
            with open(self.memory_path, "r", encoding="utf-8") as f:
                all_memory = yaml.safe_load(f) or {}
        if self.role_id in all_memory:
            self.short_momery = all_memory[self.role_id]

    def save_memory_to_file(self):
        all_memory = {}
        if os.path.exists(self.memory_path):
            with open(self.memory_path, "r", encoding="utf-8") as f:
                all_memory = yaml.safe_load(f) or {}
        all_memory[self.role_id] = self.short_momery
        with open(self.memory_path, "w", encoding="utf-8") as f:
            yaml.dump(all_memory, f, allow_unicode=True)

    async def save_memory(self, msgs, chat_history_id, executor,logger):
        if self.llm is None:
            logger.bind(tag=TAG).error("LLM is not set for memory provider")
            return None

        if len(msgs) < 2:
            return None

        msgStr = ""
        for msg in msgs:
            if msg.role == "user":
                msgStr += f"User: {msg.content}\n"
            elif msg.role == "assistant":
                msgStr += f"Assistant: {msg.content}\n"
        if len(self.short_momery) > 0:
            msgStr += "历史记忆：\n"
            msgStr += self.short_momery

        # 当前时间
        time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        msgStr += f"当前时间：{time_str}"

        result = self.llm.response_no_stream(short_term_memory_prompt, msgStr)

        json_str = extract_json_data(result)
        threading.Thread(target=inner_send_mem, args=(None, chat_history_id, json_str, logger), daemon=False).start()

        try:
            json_data = json.loads(json_str)  # 检查json格式是否正确
            self.short_momery = json_str
        except Exception as e:
            print("Error:", e)

        self.save_memory_to_file()
        logger.bind(tag=TAG).info(f"Save memory successful - Role: {self.role_id}")

        return self.short_momery

    async def query_memory(self, query: str) -> str:
        return self.short_momery
